#!/usr/bin/env python
# encoding: utf-8
'''
@file: Synthesize_more_realistic_data.py
@author: wangjianong
@time: 2020/3/18 10:48
@desc:Synthesize more realistic data
'''

import cv2
import os
import numpy as np
import random
from copy import deepcopy




def filter_smoke_data(img,thresh=15):
    '''
    过滤烟雾图片中黑色块与过小的烟雾
    :return:
    '''
    alpha = img[...,-1]
    h,w = alpha.shape
    contours,_ = cv2.findContours(alpha.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    roi_xmin = w
    roi_xmax = 0
    roi_ymin = h
    roi_ymax = 0
    for i in range(len(contours)-1,-1,-1):
        rect = cv2.boundingRect(contours[i])
        if max(rect[2],rect[3]) >= thresh:
            roi_xmin = min(roi_xmin,rect[0])
            roi_xmax = max(roi_xmax,rect[0]+rect[2])
            roi_ymin = min(roi_ymin,rect[1])
            roi_ymax = max(roi_ymax,rect[1]+rect[3])
    roi = img[roi_ymin:roi_ymax,roi_xmin:roi_xmax]
    return roi

def merge_bgfg(bg_folder,fg_folder,ratio_size=None,ratio_alpha=None):
    '''
    合成烟雾图片
    :param bg:
    :param fg:
    :param ratio:
    :return:
    '''
    #1.处理背景图片，如果size大于2000，则2倍下采样
    bg_img = cv2.imread(bg_folder, -1)
    h_bg, w_bg, c_bg = bg_img.shape
    assert c_bg == 4 or c_bg == 3, f'{bg_folder} channels={c_bg}!'
    while h_bg >= 2000 or w_bg >= 2000:
        bg_img = cv2.resize(bg_img, (w_bg // 2, h_bg // 2))
        h_bg, w_bg, c_bg = bg_img.shape

    #2.处理烟雾图片
    fg_img = cv2.imread(fg_folder, -1)
    # fg_img = filter_smoke_data(fg_img)
    h_fg, w_fg, c_fg = fg_img.shape
    if h_fg==0 or w_fg==0:
        return None,None
    assert c_fg == 4, f'{fg_folder} channels={c_fg}!'
    if ratio_size is None:
        new_h,new_w = h_fg, w_fg
    else:
        new_h,new_w = int(h_fg*ratio_size[0]),int(w_fg*ratio_size[1])
    fg_img = cv2.resize(fg_img,(new_w,new_h))
    h_fg, w_fg, c_fg = fg_img.shape

    #3.合成数据
    #3.1烟雾左上角和右下角坐标
    # tl_x = random.randint(0,w_bg-w_fg)
    # tl_y = random.randint(0,h_bg-h_fg)
    tl_x = 0
    tl_y = 0
    br_x = tl_x + w_fg
    br_y = tl_y + h_fg

    #3.2合成图片
    b, g, r, a = cv2.split(fg_img)

    ratio = 1
    # res_img = np.zeros_like(bg_img)
    res_img = deepcopy(bg_img)
    for i in range(3):
        res_img[...,i][tl_y:br_y, tl_x:br_x] = bg_img[...,i][tl_y:br_y, tl_x:br_x] * (255.0 - ratio * a) / 255
        res_img[...,i][tl_y:br_y, tl_x:br_x] += np.array(fg_img[...,i] * (ratio * a / 255), dtype=np.uint8)

    #3.3烟雾轮廓
    contours,_ = cv2.findContours(a.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        contour[...,0] += tl_x
        contour[...,1] += tl_y

    return res_img,contours

def load_smokedata(smoke_folder='/dataset/smoke_data/2.24'):
    '''
    读取烟雾图片数据
    :param smoke_folder: 烟雾数据根目录
    :return: 返回所有烟雾图片路径对应list
    '''
    dirs = [os.path.join(smoke_folder, dir) for dir in os.listdir(smoke_folder) if not dir.startswith('2.2')]
    list_smoke = []
    for dir in dirs:
        list_bgs = [os.path.join(dir,f) for f in os.listdir(dir)]
        list_smoke += list_bgs
    print(f"smoke data num:{len(list_smoke)}")
    return list_smoke

def load_backgrounds(bg_folder = '/dataset/smoke_data/all_cams'):
    '''
    读取摄像头背景数据
    :param bg_folder:
    :return:
    '''

    list_files = [os.path.join(bg_folder,f) for f in os.listdir(bg_folder)]
    return list_files

if __name__ == '__main__':
    # img_save_folder = '/dataset/smoke_data/synthetic_0318/images'
    # label_save_folder = '/dataset/smoke_data/synthetic_0318/labels'
    # os.makedirs(img_save_folder, exist_ok=True)
    # os.makedirs(label_save_folder, exist_ok=True)
    #加载烟雾数据对应文件路径
    list_smoke = load_smokedata('G:\\2.24\\2.24/')
    #加载背景数据对应路径
    list_files = load_backgrounds('G:\\1\\3\\4\images\images/')

    save_id = 0
    bg_folder = random.sample(list_files, 1)[0]
    for i in range(len(list_smoke)):
        # bg_folder = random.sample(list_files, 1)[0]
        smoke_folder = list_smoke[i]
        # ratio_size = (random.uniform(0.8,1.2),random.uniform(0.8,1.2))
        ratio_size = None
        ratio_alpha = None
        # ratio_alpha = (random.uniform(0.5,1.),random.uniform(0.5,1.))
        res_img,contours = merge_bgfg(bg_folder,smoke_folder,ratio_size=ratio_size,ratio_alpha=ratio_alpha)
        if res_img is None:
            continue
        cv2.imshow('img',res_img)
        fg_img = cv2.imread(smoke_folder,-1)
        cv2.imshow('smoke',fg_img)
        cv2.waitKey(1)

        save_id += 1






